LFM2 #20797

paulpak58 · 2025-07-11T05:23:45Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Test Plan

Test Result

(Optional) Documentation Update

…with Mamba2

…ormers >= 4.54.0.dev0

…usal_conv1d kernel

…n 4.54.0.dev0

github-actions · 2025-07-11T05:23:54Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

gemini-code-assist

Summary of Changes

Hello @paulpak58, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces the LFM2 model architecture into vLLM, including necessary components for hybrid attention and convolution-based layers. It also incorporates changes to handle CUDA versions, transformers library updates, and integration with the vLLM v1 framework.

Highlights

CMakeLists.txt: Added logic to ignore nvToolsExt for CUDA 12.9 by creating an imported interface library if it's not already a target.
vllm/config.py: Modified get_num_layers_by_block_type to handle hybrid models, specifically addressing an attribute change in transformers library versions >= 4.54.0.dev0. It now checks for layer_types in hf_text_config if layers_block_type is not found in hf_config. Also, the logic for summing the block types has been updated to account for 'full_attention' blocks when the block type is 'attention'.
vllm/model_executor/layers/conv.py: Introduced a new ShortConv custom operator for implementing a short convolution layer, including forward passes for both native and CUDA execution. The CUDA forward pass includes logic for prefill and decode stages, utilizing causal_conv1d_fn and causal_conv1d_update respectively. The class also defines a get_state_shape method.
vllm/model_executor/models/conv_cache.py: Introduced ConvCacheParams dataclass and ConvCacheManager class to manage the convolution state cache. The ConvCacheManager inherits from ConstantSizeCache and provides methods for copying cache, retrieving tensors for the current run, and providing CUDA graph capture inputs.
vllm/model_executor/models/lfm2.py: Added new modules and classes for the LFM2 model architecture, including LFM2MLP, LFM2Attention, LFM2AttentionDecoderLayer, LFM2ShortConvDecoderLayer, LFM2Model, and LFM2ForCausalLM. These components define the layers and overall structure of the LFM2 model, supporting hybrid attention and convolution-based layers. The LFM2ForCausalLM class integrates the LFM2 model with the vLLM framework, including cache management and logits processing.
vllm/model_executor/models/registry.py: Registered LFM2ForCausalLM in the model registry, associating it with the 'lfm2' identifier.
vllm/transformers_utils/configs/ovis.py: Wrapped the AutoConfig.register call for AIMv2Config in a try-except block to prevent errors if AutoConfig is not available.
vllm/utils/init.py: Added conv to the LayerBlockType enum.
vllm/v1/attention/backends/mamba_attn.py: Added get_short_conv_chunk_size function and updated Mamba2AttentionMetadataBuilder to support ShortConvSpec.
vllm/v1/core/single_type_kv_cache_manager.py: Added ShortConvSpec to the dictionary of kv cache specs.
vllm/v1/kv_cache_interface.py: Added ShortConvSpec dataclass to define the specification for the short convolution KV cache.
vllm/v1/worker/gpu_model_runner.py: Imported ShortConv, updated initialize_attn_backend to support ShortConvSpec, and modified _reshape_kv_cache_tensors and get_kv_cache_spec to handle short convolution layers.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces support for the LFM2 model, a new hybrid architecture. The changes are comprehensive, affecting model definition, caching mechanisms, and the core engine logic. The implementation cleverly reuses existing infrastructure for Mamba-like layers to handle the new convolution layers, which is a great approach.

I've provided several suggestions to enhance code clarity, maintainability, and robustness. These include refactoring duplicated code, simplifying complex conditions, and improving exception handling. Overall, this is a solid contribution.

vllm/transformers_utils/configs/ovis.py

gemini-code-assist · 2025-07-11T05:27:30Z

vllm/config.py

+                return sum(
+                    1 for t in layers_block_type_value[start:end]
+                    if (t == "full_attention" and "attention" == block_type.value) or (t == block_type.value)
+                )


This conditional logic for counting layer types is a bit complex and hard to read. It can be simplified by handling the special case for attention layers separately, which would make the code more readable and easier to maintain.

Suggested change

return sum(

1 for t in layers_block_type_value[start:end]

if (t == "full_attention" and "attention" == block_type.value) or (t == block_type.value)

)

if block_type == LayerBlockType.attention:

return sum(t in ("attention", "full_attention")

for t in layers_block_type_value[start:end])

return sum(t == block_type.value

for t in layers_block_type_value[start:end])

vllm/model_executor/layers/conv.py

gemini-code-assist · 2025-07-11T05:27:30Z

vllm/model_executor/layers/conv.py

+        return contextualized_states
+
+
+    def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:


The return type hint tuple[tuple[int, ...], tuple[int, ...]] indicates a tuple containing two tuples of integers. However, the function returns a tuple containing only one tuple: (conv_state_shape,).

To match the implementation and the expected usage with MambaSpec-like structures, the type hint should be tuple[tuple[int, ...], ...], which correctly represents a tuple containing one or more tuples of integers.

Suggested change

def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:

def get_state_shape(self) -> tuple[tuple[int, ...], ...]:

vllm/v1/kv_cache_interface.py

vllm/v1/worker/gpu_model_runner.py

mergify · 2025-07-11T05:56:40Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @paulpak58.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Paul Pak <paulpak58@gmail.com>

paulpak58 · 2025-07-15T21:48:30Z

CMakeLists.txt

+if (NOT TARGET CUDA::nvToolsExt)
+    add_library(CUDA::nvToolsExt INTERFACE IMPORTED)
+endif()
+


possibly a cleaner solution than this, but this works.

mergify · 2025-07-16T02:26:06Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @paulpak58.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Paul Pak <paulpak58@gmail.com>

paulpak58 added 7 commits July 10, 2025 19:44

[cmake] ignore nvToolsExt for cuda-12.9

3a50223

[model_executor][models] LFM2 architecture

6fd86d2

[configs] use layer_types from huggingface hybrids >= 4.54.0.dev0

aaf7df1

[model_runner][v1] ShortConvSpec for ShortConv layers; compatibility …

6c80caf

…with Mamba2

[configs] need to detect full_attention key in layer_types for transf…

d17c95f

…ormers >= 4.54.0.dev0

[layers][conv] update ShortConv layer to be compatible with triton ca…

1bc8835

…usal_conv1d kernel

[transformers][ovis] tmp: AIMv2Config doesn't need to be registered o…

e550362

…n 4.54.0.dev0

mergify bot added ci/build new-model Requests to new models v1 labels Jul 11, 2025

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

paulpak58 changed the title ~~Lfm2~~ LFM2 Jul 11, 2025

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

mergify bot added the needs-rebase label Jul 11, 2025

paulpak58 added 4 commits July 11, 2025 16:17

[models][lfm2] LFM2->Lfm2 to match config

05af65a

[merge] upstream @ 5bac613

40d81e9

[v1][cache] add support for conv cache shapes

7241660

[v1][config] generalize HybridAttentionMambaModelConfig

3d3be6a

mergify bot removed the needs-rebase label Jul 15, 2025

paulpak58 added 3 commits July 15, 2025 21:22

[merge] upstream main @ 19c8630

b2447dd

[layer][conv] update conv metadata in causal_conv1d

46902dc

[misc] format + cleanup

260e3fe

Signed-off-by: Paul Pak <paulpak58@gmail.com>

paulpak58 commented Jul 15, 2025

View reviewed changes

mergify bot added the needs-rebase label Jul 16, 2025

paulpak58 added 2 commits July 16, 2025 15:29

[layers][conv] fix minor discprencies in decode conv

1dff6e1

Signed-off-by: Paul Pak <paulpak58@gmail.com>

[merge] upstream @ a0f8a79

30621b4

mergify bot removed the needs-rebase label Jul 16, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

LFM2 #20797

LFM2 #20797

paulpak58 commented Jul 11, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jul 11, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

gemini-code-assist bot Jul 11, 2025

Uh oh!

Uh oh!

gemini-code-assist bot Jul 11, 2025

Uh oh!

Uh oh!

Uh oh!

mergify bot commented Jul 11, 2025

Uh oh!

paulpak58 Jul 15, 2025

Uh oh!

mergify bot commented Jul 16, 2025

Uh oh!

Uh oh!

-                return sum(
-for t in layers_block_type_value[start:end]
-                    if (t == "full_attention" and "attention" == block_type.value) or (t == block_type.value)
-                )
+                if block_type == LayerBlockType.attention:
+                    return sum(t in ("attention", "full_attention")
+                               for t in layers_block_type_value[start:end])
+                return sum(t == block_type.value
+                           for t in layers_block_type_value[start:end])

		return contextualized_states


		def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:

Uh oh!

LFM2 #20797

Are you sure you want to change the base?

LFM2 #20797

Conversation

paulpak58 commented Jul 11, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jul 11, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

mergify bot commented Jul 11, 2025

Uh oh!

paulpak58 Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Jul 16, 2025

Uh oh!

Uh oh!

paulpak58 commented Jul 11, 2025 •

edited by github-actions bot

Loading